from agent.Base_Agent import Base_Agent
from math_ops.Math_Ops import Math_Ops as M
import math
import numpy as np


class Agent(Base_Agent):
    
    def __init__(self, host:str, agent_port:int, monitor_port:int, unum:int,
                 team_name:str, enable_log, enable_draw, wait_for_server=True, is_fat_proxy=False) -> None:
        
        # define robot type
        robot_type = (0,1,1,1,2,3,3,3,4,4,4)[unum-1]

        # Initialize base agent
        # Args: Server IP, Agent Port, Monitor Port, Uniform No., Robot Type, Team Name, Enable Log, Enable Draw, play mode correction, Wait for Server, Hear Callback
        super().__init__(host, agent_port, monitor_port, unum, robot_type, team_name, enable_log, enable_draw, True, wait_for_server, None)

        self.enable_draw = enable_draw
        self.state = 0  # 0-Normal, 1-Getting up, 2-Kicking
        self.kick_direction = 0
        self.kick_distance = 0
        self.fat_proxy_cmd = "" if is_fat_proxy else None
        self.fat_proxy_walk = np.zeros(3) # filtered walk parameters for fat proxy
# DWJ-----------------------------------------------------------------------------------------------------------------
        self.ball_3s_speed_pos = (0, 0, 0)
# DWJ -------------------------------------------------------------------------------------------------------------------

        self.init_pos = ([-14.5,0],[-9,-5],[-9,0],[-9,5],[-5,-5],[-5,0],[-5,5],[-1,-6],[-1,-2.5],[-1,2.5],[-1,6])[unum-1] # initial formation


    def beam(self, avoid_center_circle=False):
        r = self.world.robot
        pos = self.init_pos[:] # copy position list 
        self.state = 0

        # Avoid center circle by moving the player back 
        if avoid_center_circle and np.linalg.norm(self.init_pos) < 2.5:
            pos[0] = -2.3 

        if np.linalg.norm(pos - r.loc_head_position[:2]) > 0.1 or self.behavior.is_ready("Get_Up"):
            self.scom.commit_beam(pos, M.vector_angle((-pos[0],-pos[1]))) # beam to initial position, face coordinate (0,0)
        else:
            if self.fat_proxy_cmd is None: # normal behavior
                self.behavior.execute("Zero_Bent_Knees_Auto_Head")
            else: # fat proxy behavior
                self.fat_proxy_cmd += "(proxy dash 0 0 0)"
                self.fat_proxy_walk = np.zeros(3) # reset fat proxy walk


    # 移动
    def move(self, target_2d=(0,0), orientation=None, is_orientation_absolute=True,
             avoid_obstacles=True, priority_unums=[], is_aggressive=False, timeout=3000):
        '''
        Walk to target position

        Parameters
        ----------
        target_2d : array_like
            2D target in absolute coordinates
        orientation : float
            absolute or relative orientation of torso, in degrees
            set to None to go towards the target (is_orientation_absolute is ignored)
        is_orientation_absolute : bool
            True if orientation is relative to the field, False if relative to the robot's torso
        avoid_obstacles : bool
            True to avoid obstacles using path planning (maybe reduce timeout arg if this function is called multiple times per simulation cycle)
        priority_unums : list
            list of teammates to avoid (since their role is more important)
        is_aggressive : bool
            if True, safety margins are reduced for opponents
        timeout : float
            restrict path planning to a maximum duration (in microseconds)    
        '''
        r = self.world.robot

        if self.fat_proxy_cmd is not None: # fat proxy behavior
            self.fat_proxy_move(target_2d, orientation, is_orientation_absolute) # ignore obstacles
            return

        if avoid_obstacles:
            target_2d, _, distance_to_final_target = self.path_manager.get_path_to_target(
                target_2d, priority_unums=priority_unums, is_aggressive=is_aggressive, timeout=timeout)
        else:
            distance_to_final_target = np.linalg.norm(target_2d - r.loc_head_position[:2])

        self.behavior.execute("Walk", target_2d, True, orientation, is_orientation_absolute, distance_to_final_target) # Args: target, is_target_abs, ori, is_ori_abs, distance





    def kick(self, kick_direction=None, kick_distance=None, abort=False, enable_pass_command=False):
        '''
        Walk to ball and kick

        Parameters
        ----------
        kick_direction : float
            kick direction, in degrees, relative to the field
        kick_distance : float
            kick distance in meters
        abort : bool
            True to abort.
            The method returns True upon successful abortion, which is immediate while the robot is aligning itself. 
            However, if the abortion is requested during the kick, it is delayed until the kick is completed.
        avoid_pass_command : bool
            When False, the pass command will be used when at least one opponent is near the ball
            
        Returns
        -------
        finished : bool
            Returns True if the behavior finished or was successfully aborted.
        '''

        if self.min_opponent_ball_dist < 1.45 and enable_pass_command:
            self.scom.commit_pass_command()

        self.kick_direction = self.kick_direction if kick_direction is None else kick_direction
        self.kick_distance = self.kick_distance if kick_distance is None else kick_distance

        if self.fat_proxy_cmd is None: # normal behavior
            return self.behavior.execute("Basic_Kick", self.kick_direction, abort) # Basic_Kick has no kick distance control
        else: # fat proxy behavior
            return self.fat_proxy_kick()




    def think_and_send(self):
        # ========================= 基础数据准备 ========================= 
        # 获取世界状态、机器人对象、当前时间等核心数据 
        w = self.world   # 世界状态对象 
        r = self.world.robot   # 当前机器人对象 
        my_head_pos_2d = r.loc_head_position[:2]   # 机器人头部二维坐标（x,y）
        my_ori = r.imu_torso_orientation   # 机器人躯干朝向角度 
        ball_2d = w.ball_abs_pos[:2]   # 足球绝对二维坐标 
        ball_vec = ball_2d - my_head_pos_2d  # 球相对于头部的向量 
        ball_dir = M.vector_angle(ball_vec)   # 球相对于头部的方向角（弧度）
        ball_dist = np.linalg.norm(ball_vec)   # 球与头部的直线距离 
        ball_sq_dist = ball_dist * ball_dist  # 平方距离（用于快速比较）
        ball_speed = np.linalg.norm(w.get_ball_abs_vel(6)[:2])   # 球速（取最近6ms速度）
        behavior = self.behavior   # 行为控制模块 
        goal_dir = M.target_abs_angle(ball_2d,  (15.05,0))  # 球到对方球门的绝对角度 
        path_draw_options = self.path_manager.draw_options   # 路径绘制参数 
        PM = w.play_mode   # 当前比赛模式 
        PM_GROUP = w.play_mode_group   # 比赛模式分组 
    
        # ========================= 预处理计算 ========================= 
        # 预测球速降低到0.5m/s时的未来位置（用于策略决策）
        slow_ball_pos = w.get_predicted_ball_pos(0.5) 
        
        # 计算队友到球的平方距离（排除无效队友）
        teammates_ball_sq_dist = [
            np.sum((p.state_abs_pos[:2]  - slow_ball_pos) ** 2)  # 计算平方距离 
            if p.state_last_update  != 0  # 排除无状态更新的队友 
            and (w.time_local_ms  - p.state_last_update  <= 360 or p.is_self)   # 状态有效性检查（360ms内）
            and not p.state_fallen   # 排除跌倒的队友 
            else 1000  # 无效情况设为最大距离 
            for p in w.teammates  
        ]
        
        # 同理计算对手到球距离（逻辑与队友计算相同）
        opponents_ball_sq_dist = [...]  # 此处省略重复逻辑 
        
        # 获取最小距离和对应球员编号 
        min_teammate_ball_sq_dist = min(teammates_ball_sq_dist)
        self.min_teammate_ball_dist  = math.sqrt(min_teammate_ball_sq_dist)   # 转化为实际距离 
        self.min_opponent_ball_dist  = math.sqrt(min(opponents_ball_sq_dist)) 
        active_player_unum = teammates_ball_sq_dist.index(min_teammate_ball_sq_dist)  + 1  # 获取最近队友编号 
    
        # ========================= 决策逻辑 ========================= 
        if PM == w.M_GAME_OVER:  # 比赛结束状态 
            pass  # 无动作 
        elif PM_GROUP == w.MG_ACTIVE_BEAM:  # 主动定位球状态 
            self.beam()   # 执行定位球动作 
        elif PM_GROUP == w.MG_PASSIVE_BEAM:
            self.beam(True)   # 被动定位球（避开中圈）
        elif self.state  == 1 or (behavior.is_ready("Get_Up")  and self.fat_proxy_cmd  is None):
            self.state  = 0 if behavior.execute("Get_Up")  else 1  # 摔倒后起身逻辑 
        elif PM == w.M_OUR_KICKOFF:  # 我方开球 
            if r.unum  == 9:  # 特定球员执行开球 
                self.kick(120,  3)  # 执行强力踢球（角度120度，力量3）
            else:
                self.move(self.init_pos,  orientation=ball_dir)  # 其他球员保持原位 
        elif active_player_unum != r.unum:   # 非当前活跃球员 
            if r.unum  == 1:  # 守门员逻辑 
                if ball_dist < 0.5 and ball_speed < 0.1:
                    self.scom.commit_pass_command()
                    pass
                else:
                    self.move(self.init_pos,  orientation=ball_dir)  # 保持门线位置 
            else:  # 普通球员站位逻辑 
                # 动态调整阵型位置（基于球的位置）
                new_x = max(0.5, (ball_2d[0]+15)/15) * (self.init_pos[0]+15)  - 15 
                if self.min_teammate_ball_dist  < self.min_opponent_ball_dist: 
                    new_x = min(new_x + 3.5, 13)  # 团队控球时前压 
                self.move((new_x,  self.init_pos[1]),  orientation=ball_dir, priority_unums=[active_player_unum])
    
        else:  # 当前是活跃球员 
            # 启用路径绘制（用于调试）
            path_draw_options(enable_obstacles=True, enable_path=True, use_team_drawing_channel=True)
            
            # 进攻逻辑（比赛进行中且球在前场时允许传球）
            enable_pass_command = (PM == w.M_PLAY_ON and ball_2d[0]<6)
            
            if PM == w.M_OUR_CORNER_KICK:  # 角球战术 
                self.kick(-np.sign(ball_2d[1])*95,  5.5)  # 斜传禁区 
            elif self.min_opponent_ball_dist  + 0.5 < self.min_teammate_ball_dist:   # 防守逻辑 
                if self.state  == 2:  # 中断当前动作 
                    self.state  = 0 if self.kick(abort=True)  else 2 
                else:  # 卡位防守（站在球与球门连线之间）
                    self.move(slow_ball_pos  + M.normalize_vec((-16,0)  - slow_ball_pos) * 0.2, is_aggressive=True)
            else:  # 常规进攻 
                self.state  = 0 if self.kick(goal_dir,  9, False, enable_pass_command) else 2  # 朝球门方向射门 
    
            path_draw_options(enable_obstacles=False, enable_path=False)  # 关闭调试绘制 
    
        # ========================= 指令传输 ========================= 
        self.radio.broadcast()   # 发送无线通信 
        if self.fat_proxy_cmd  is None:  # 正常指令模式 
            self.scom.commit_and_send(r.get_command())   # 发送动作指令 
        else:  # 紧急指令模式 
            self.scom.commit_and_send(self.fat_proxy_cmd.encode()) 
            self.fat_proxy_cmd  = ""
    
        # ========================= 调试辅助 ========================= 
        if self.enable_draw:   # 可视化调试信息 
            d = w.draw  
            if active_player_unum == r.unum:   # 仅活跃球员显示预测轨迹 
                d.point(slow_ball_pos,  3, d.Color.pink,  "status", False)  # 显示预测球位 
                d.point(w.ball_2d_pred_pos[-1],  5, d.Color.pink,  "status", False)  # 最终预测位置 
                d.annotation((*my_head_pos_2d,  0.6), "I've got it!", d.Color.yellow,  "status")  # 状态提示 
            else:
                d.clear("status")   # 非活跃球员清空绘图 




    #--------------------------------------- Fat proxy auxiliary methods


    def fat_proxy_kick(self):
        w = self.world
        r = self.world.robot 
        ball_2d = w.ball_abs_pos[:2]
        my_head_pos_2d = r.loc_head_position[:2]

        if np.linalg.norm(ball_2d - my_head_pos_2d) < 0.25:
            # fat proxy kick arguments: power [0,10]; relative horizontal angle [-180,180]; vertical angle [0,70]
            self.fat_proxy_cmd += f"(proxy kick 10 {M.normalize_deg( self.kick_direction  - r.imu_torso_orientation ):.2f} 20)" 
            self.fat_proxy_walk = np.zeros(3) # reset fat proxy walk
            return True
        else:
            self.fat_proxy_move(ball_2d-(-0.1,0), None, True) # ignore obstacles
            return False


    def fat_proxy_move(self, target_2d, orientation, is_orientation_absolute):
        r = self.world.robot

        target_dist = np.linalg.norm(target_2d - r.loc_head_position[:2])
        target_dir = M.target_rel_angle(r.loc_head_position[:2], r.imu_torso_orientation, target_2d)

        if target_dist > 0.1 and abs(target_dir) < 8:
            self.fat_proxy_cmd += (f"(proxy dash {100} {0} {0})")
            return

        if target_dist < 0.1:
            if is_orientation_absolute:
                orientation = M.normalize_deg( orientation - r.imu_torso_orientation )
            target_dir = np.clip(orientation, -60, 60)
            self.fat_proxy_cmd += (f"(proxy dash {0} {0} {target_dir:.1f})")
        else:
            self.fat_proxy_cmd += (f"(proxy dash {20} {0} {target_dir:.1f})")